"""
投资团队工作流
定义投资团队成员协作的工作流程
"""
from typing import Dict, List, Any, Optional, Callable
from loguru import logger
import time
import datetime
from pathlib import Path

from crewai import Crew, Process, Task
from langgraph.graph import StateGraph
# 使用纯字典格式工具，避免Tool类导入

from agents.base_agent import BaseInvestmentAgent
from agents.analyst_agent import MarketAnalyst, StockAnalyst
from agents.strategist_agent import StrategyAdvisor, PortfolioManager
from agents.decision_agent import ChiefInvestmentOfficer, ResearchSpecialist
from database.data_collector import DataCollector
from utils.stock_analysis import calculate_technical_indicators, analyze_financial_data
from utils.portfolio_management import optimize_portfolio, calculate_risk_metrics
from workflows.base_workflow import BaseWorkflow, WorkflowState

# 定义工作流状态类型
InvestmentTeamState = Dict[str, Any]

def create_investment_team():
    """创建投资团队工作流"""
    logger.info("创建投资团队工作流")
    
    # 创建工作流
    workflow = InvestmentTeamWorkflow()
    
    return workflow


class InvestmentTeamWorkflow(BaseWorkflow):
    """投资团队工作流"""
    
    def __init__(self):
        """初始化投资团队工作流"""
        super().__init__("A股投资团队")
        
        # 初始化状态
        initial_state = {
            "date": datetime.datetime.now().strftime("%Y-%m-%d"),
            "market_analysis": {},
            "stock_analysis": {},
            "investment_strategy": {},
            "portfolio": {},
            "decisions": {},
            "execution_results": {}
        }
        
        self.state = WorkflowState(initial_state)
        
        # 初始化团队成员
        self._init_team_members()
        
        # 添加工作流步骤
        self.add_step(self._collect_market_data, "收集市场数据")
        self.add_step(self._analyze_market, "分析市场环境")
        self.add_step(self._research_stocks, "研究个股")
        self.add_step(self._develop_strategy, "制定投资策略")
        self.add_step(self._make_decisions, "做出投资决策")
        self.add_step(self._execute_trades, "执行交易")
        self.add_step(self._review_performance, "评估投资表现")
    
    def _init_team_members(self):
        """初始化团队成员"""
        logger.info("初始化投资团队成员")
        
        # 创建团队成员智能体
        self.market_analyst = MarketAnalyst(verbose=True)
        self.stock_analyst = StockAnalyst(verbose=True)
        self.strategy_advisor = StrategyAdvisor(verbose=True)
        self.portfolio_manager = PortfolioManager(verbose=True)
        self.chief_investment_officer = ChiefInvestmentOfficer(verbose=True)
        self.tech_specialist = ResearchSpecialist(specialization="科技", verbose=True)
        self.consumer_specialist = ResearchSpecialist(specialization="消费", verbose=True)
        
        # 创建数据采集器
        self.data_collector = DataCollector()
    
    def _collect_market_data(self, state: InvestmentTeamState, previous_result: Any) -> InvestmentTeamState:
        """
        收集市场数据
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        logger.info("收集市场数据")
        
        try:
            # 获取当前日期
            current_date = datetime.datetime.now().strftime("%Y%m%d")
            
            # 获取上证指数数据
            sh_index_data = self.data_collector.get_index_data("000001", "20230101", current_date)
            
            # 获取深证成指数据
            sz_index_data = self.data_collector.get_index_data("399001", "20230101", current_date)
            
            # 获取沪深300数据
            hs300_data = self.data_collector.get_index_data("000300", "20230101", current_date)
            
            # 获取股票列表
            stock_list = self.data_collector.get_stock_list()
            
            # 更新状态
            market_data = {
                "date": current_date,
                "indices": {
                    "sh_index": sh_index_data.to_dict() if not sh_index_data.empty else {},
                    "sz_index": sz_index_data.to_dict() if not sz_index_data.empty else {},
                    "hs300": hs300_data.to_dict() if not hs300_data.empty else {}
                },
                "stock_list": stock_list.to_dict() if not stock_list.empty else {}
            }
            
            # 返回更新状态
            return {"market_data": market_data}
        
        except Exception as e:
            logger.error(f"收集市场数据失败: {e}")
            return {"market_data": {}, "error": str(e)}
    
    def _analyze_market(self, state: InvestmentTeamState, previous_result: Any) -> InvestmentTeamState:
        """
        分析市场环境
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        logger.info("分析市场环境")
        
        max_retries = 3
        retry_delay = 1
        
        # 获取市场数据
        market_data = state.get("market_data", {})
        
        # 默认返回值
        default_result = {
            "market_analysis": {
                "date": state.get("date", ""),
                "market_trend": "unknown",
                "sector_rotation": {},
                "economic_indicators": {},
                "risk_assessment": {},
                "summary": "分析失败"
            },
            "error": "分析失败"
        }
        
        if not market_data:
            logger.warning("没有市场数据，无法进行分析")
            default_result["error"] = "没有市场数据"
            return default_result
                
        for attempt in range(max_retries):
            try:
                # 创建市场分析任务
                market_analysis_task = Task(
                    description="分析当前市场趋势、行业轮动和宏观经济指标",
                    expected_output="详细的市场分析报告，包括市场趋势、投资机会和风险提示",
                    agent=self.market_analyst.get_agent()
                )
                
                # 创建CrewAI团队执行分析
                analysis_crew = Crew(
                    agents=[self.market_analyst.get_agent()],
                    tasks=[market_analysis_task],
                    verbose=True,
                    process=Process.sequential
                )
                
                # 执行分析
                market_analysis_result = analysis_crew.kickoff(
                    inputs={
                        "market_data": f"市场数据概览: 上证指数、深证成指、沪深300等主要指数信息",
                        "date": state.get("date", "")
                    }
                )
                
                # 解析分析结果
                market_analysis = {
                    "date": state.get("date", ""),
                    "market_trend": "上涨",  # 这会被实际结果替换
                    "sector_rotation": {},
                    "economic_indicators": {},
                    "risk_assessment": {},
                    "market_sentiment": {
                        "fear_greed_index": 50,  # 默认中性
                        "market_breadth": 50,     # 默认中性
                        "volatility_index": 20.0, # 默认波动率
                        "north_fund": 0.0         # 默认北向资金
                    },
                    "summary": str(market_analysis_result.raw) if hasattr(market_analysis_result, "raw") else str(market_analysis_result)
                }
                
                # 返回更新状态
                return {"market_analysis": market_analysis}
                
            except Exception as e:
                logger.warning(f"市场分析尝试 {attempt + 1}/{max_retries} 失败: {str(e)}")
                if attempt < max_retries - 1:
                    time.sleep(retry_delay * (attempt + 1))
                    continue
                
                logger.error("所有市场分析尝试均失败，使用本地分析")
                try:
                    # 本地降级分析
                    from utils.market_analysis import local_analyze_market
                    market_analysis = local_analyze_market(market_data)
                    # 确保本地分析也包含market_sentiment
                    if "market_sentiment" not in market_analysis:
                        market_analysis["market_sentiment"] = {
                            "fear_greed_index": 50,
                            "market_breadth": 50,
                            "volatility_index": 20.0,
                            "north_fund": 0.0
                        }
                    return {"market_analysis": market_analysis, "warning": "使用本地分析结果"}
                except Exception as fallback_error:
                    logger.error(f"本地市场分析也失败: {fallback_error}")
                    default_result["error"] = str(fallback_error)
                    return default_result
        
        return default_result
    
    def _research_stocks(self, state: InvestmentTeamState, previous_result: Any) -> InvestmentTeamState:
        """
        研究个股
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        logger.info("研究个股")
        
        try:
            # 获取市场分析结果
            market_analysis = state.get("market_analysis", {})
            
            if not market_analysis:
                logger.warning("没有市场分析结果，无法进行个股研究")
                return {"stock_analysis": {}, "error": "没有市场分析结果"}
            
            # 基于市场分析，选择要研究的股票
            from config.settings import TARGET_STOCKS
            target_stocks = TARGET_STOCKS
            
            # 创建个股分析任务
            stock_analysis_tasks = []
            
            for stock_code in target_stocks:
                task = Task(
                    description=f"分析股票 {stock_code} 的基本面和技术面",
                    expected_output=f"股票 {stock_code} 的详细分析报告，包括投资价值评估",
                    agent=self.stock_analyst.get_agent()
                )
                stock_analysis_tasks.append(task)
            
            # 创建行业研究任务
            tech_research_task = Task(
                description="研究科技行业发展趋势和投资机会",
                expected_output="科技行业详细研究报告，包括重点公司推荐",
                agent=self.tech_specialist.get_agent()
            )
            
            consumer_research_task = Task(
                description="研究消费行业发展趋势和投资机会",
                expected_output="消费行业详细研究报告，包括重点公司推荐",
                agent=self.consumer_specialist.get_agent()
            )
            
            # 创建CrewAI团队执行研究
            research_crew = Crew(
                agents=[
                    self.stock_analyst.get_agent(),
                    self.tech_specialist.get_agent(),
                    self.consumer_specialist.get_agent()
                ],
                tasks=stock_analysis_tasks + [tech_research_task, consumer_research_task],
                verbose=True,
                process=Process.sequential
            )
            
            # 执行研究
            # 确保传入的是字符串，而不是CrewOutput对象
            if market_analysis.get("summary") is not None:
                if hasattr(market_analysis.get("summary"), "raw"):
                    # 如果是CrewOutput对象，提取raw属性
                    market_analysis_str = str(market_analysis.get("summary").raw)
                else:
                    # 如果已经是字符串或其他可序列化类型
                    market_analysis_str = str(market_analysis.get("summary"))
            else:
                market_analysis_str = ""
            
            stock_research_result = research_crew.kickoff(
                inputs={
                    "market_analysis": market_analysis_str,
                    "date": state.get("date", "")
                }
            )
            
            # 解析研究结果
            stock_analysis = {
                "date": state.get("date", ""),
                "analyzed_stocks": {
                    stock_code: calculate_technical_indicators(stock_code)
                    for stock_code in target_stocks
                },
                "industry_research": {
                    "tech": "科技行业研究结果",
                    "consumer": "消费行业研究结果"
                },
                "investment_ideas": [],
                "summary": str(stock_research_result.raw) if hasattr(stock_research_result, "raw") else str(stock_research_result)
            }
            
            # 返回更新状态
            return {"stock_analysis": stock_analysis}
        
        except Exception as e:
            logger.error(f"研究个股失败: {e}")
            return {"stock_analysis": {}, "error": str(e)}
    
    def _develop_strategy(self, state: InvestmentTeamState, previous_result: Any) -> InvestmentTeamState:
        """
        制定投资策略
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        from config.settings import SECTOR_ALLOCATION
        """
        制定投资策略
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        logger.info("制定投资策略")
        
        try:
            # 获取市场分析和个股研究结果
            market_analysis = state.get("market_analysis", {})
            stock_analysis = state.get("stock_analysis", {})
            
            if not market_analysis or not stock_analysis:
                logger.warning("缺少市场分析或个股研究结果，无法制定投资策略")
                return {"investment_strategy": {}, "error": "缺少分析结果"}
            
            # 创建策略制定任务
            strategy_task = Task(
                description="基于市场分析和个股研究，制定最优投资策略",
                expected_output="详细的投资策略报告，包括资产配置和风险管理方案",
                agent=self.strategy_advisor.get_agent()
            )
            
            # 创建CrewAI团队执行策略制定
            strategy_crew = Crew(
                agents=[self.strategy_advisor.get_agent()],
                tasks=[strategy_task],
                verbose=True,
                process=Process.sequential
            )
            
            # 执行策略制定
            # 转换CrewOutput为字符串
            market_analysis_str = str(market_analysis.get("summary", ""))
            stock_analysis_str = str(stock_analysis.get("summary", ""))
            
            strategy_result = strategy_crew.kickoff(
                inputs={
                    "market_analysis": market_analysis_str,
                    "stock_analysis": stock_analysis_str,
                    "date": state.get("date", "")
                }
            )
            
            # 解析策略结果
            investment_strategy = {
                "date": state.get("date", ""),
                "market_view": "看多",  # 这会被实际结果替换
                "asset_allocation": {
                    "stocks": 0.7,
                    "bonds": 0.2,
                    "cash": 0.1
                },
                # 从配置中获取行业配置
                "sector_allocation": SECTOR_ALLOCATION,
                "risk_management": {
                    "stop_loss": 0.1,
                    "position_sizing": "moderate",
                    "diversification": "high"
                },
                "summary": str(strategy_result.raw) if hasattr(strategy_result, "raw") else str(strategy_result)
            }
            
            # 优化资产配置
            from config.settings import TARGET_STOCKS
            portfolio_data = {
                "target_stocks": TARGET_STOCKS,
                "risk_preference": "balanced",
                "current_weights": {}
            }
            
            optimized_weights = optimize_portfolio(portfolio_data)
            investment_strategy["optimized_weights"] = optimized_weights
            
            # 返回更新状态
            return {"investment_strategy": investment_strategy}
        
        except Exception as e:
            logger.error(f"制定投资策略失败: {e}")
            return {"investment_strategy": {}, "error": str(e)}
    
    def _make_decisions(self, state: InvestmentTeamState, previous_result: Any) -> InvestmentTeamState:
        """
        做出投资决策
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        logger.info("做出投资决策")
        
        try:
            # 获取投资策略
            investment_strategy = state.get("investment_strategy", {})
            market_analysis = state.get("market_analysis", {})
            stock_analysis = state.get("stock_analysis", {})
            
            if not investment_strategy:
                logger.warning("没有投资策略，无法做出投资决策")
                return {"decisions": {}, "error": "没有投资策略"}
            
            # 创建决策任务
            decision_task = Task(
                description="基于投资策略，做出具体投资决策",
                expected_output="详细的投资决策报告，包括具体买入卖出建议",
                agent=self.chief_investment_officer.get_agent()
            )
            
            # 创建CrewAI团队执行决策
            decision_crew = Crew(
                agents=[self.chief_investment_officer.get_agent()],
                tasks=[decision_task],
                verbose=True,
                process=Process.sequential
            )
            
            # 执行决策
            # 转换CrewOutput为字符串
            investment_strategy_str = str(investment_strategy.get("summary", ""))
            market_analysis_str = str(market_analysis.get("summary", ""))
            stock_analysis_str = str(stock_analysis.get("summary", ""))
            
            decision_result = decision_crew.kickoff(
                inputs={
                    "investment_strategy": investment_strategy_str,
                    "market_analysis": market_analysis_str,
                    "stock_analysis": stock_analysis_str,
                    "date": state.get("date", "")
                }
            )
            
            # 解析决策结果
            decisions = {
                "date": state.get("date", ""),
                "trades": [],
                "portfolio_adjustments": {
                    "increase_sectors": ["消费", "医疗"],
                    "decrease_sectors": ["金融", "地产"]
                },
                "risk_controls": {
                    "stop_loss_orders": [
                        {"stock_code": "600519", "price": 1600},
                        {"stock_code": "000858", "price": 58}
                    ]
                },
                "summary": str(decision_result.raw) if hasattr(decision_result, "raw") else str(decision_result)
            }
            
            # 返回更新状态
            return {"decisions": decisions}
        
        except Exception as e:
            logger.error(f"做出投资决策失败: {e}")
            return {"decisions": {}, "error": str(e)}
    
    def _execute_trades(self, state: InvestmentTeamState, previous_result: Any) -> InvestmentTeamState:
        """
        执行交易
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        logger.info("执行交易")
        
        try:
            # 获取投资决策
            decisions = state.get("decisions", {})
            
            if not decisions:
                logger.warning("没有投资决策，无法执行交易")
                return {"execution_results": {}, "error": "没有投资决策"}
            
            # 创建交易任务
            execution_task = Task(
                description="执行投资决策中的交易指令",
                expected_output="交易执行结果报告，包括成交价格和数量",
                agent=self.portfolio_manager.get_agent()
            )
            
            # 创建CrewAI团队执行交易
            execution_crew = Crew(
                agents=[self.portfolio_manager.get_agent()],
                tasks=[execution_task],
                verbose=True,
                process=Process.sequential
            )
            
            # 执行交易
            # 转换决策为字符串
            decisions_str = str(decisions)
            
            execution_result = execution_crew.kickoff(
                inputs={
                    "decisions": decisions_str
                }
            )
            
            # 模拟交易执行结果
            trades = decisions.get("trades", [])
            executed_trades = []
            
            for trade in trades:
                executed_trade = trade.copy()
                executed_trade["status"] = "executed"
                executed_trade["executed_price"] = trade.get("price", 0)
                executed_trade["executed_amount"] = trade.get("amount", 0)
                executed_trade["execution_time"] = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                executed_trades.append(executed_trade)
            
            # 解析执行结果
            execution_results = {
                "date": state.get("date", ""),
                "executed_trades": executed_trades,
                "success_rate": 1.0,
                "average_slippage": 0.001,
                "execution_costs": 0.0005,
                "summary": str(execution_result.raw) if hasattr(execution_result, "raw") else str(execution_result)
            }
            
            # 返回更新状态
            return {"execution_results": execution_results}
        
        except Exception as e:
            logger.error(f"执行交易失败: {e}")
            return {"execution_results": {}, "error": str(e)}
    
    def _review_performance(self, state: InvestmentTeamState, previous_result: Any) -> InvestmentTeamState:
        """
        评估投资表现
        
        参数:
            state: 工作流状态
            previous_result: 上一步骤的结果
            
        返回:
            更新后的状态
        """
        logger.info("评估投资表现")
        
        try:
            # 获取执行结果
            execution_results = state.get("execution_results", {})
            investment_strategy = state.get("investment_strategy", {})
            
            if not execution_results:
                logger.warning("没有执行结果，无法评估投资表现")
                return {"performance": {}, "error": "没有执行结果"}
            
            # 创建评估任务
            review_task = Task(
                description="评估投资决策和执行的表现",
                expected_output="详细的投资表现评估报告，包括收益分析和改进建议",
                agent=self.chief_investment_officer.get_agent()
            )
            
            # 创建CrewAI团队执行评估
            review_crew = Crew(
                agents=[self.chief_investment_officer.get_agent()],
                tasks=[review_task],
                verbose=True,
                process=Process.sequential
            )
            
            # 执行评估
            # 转换输入为字符串
            execution_results_str = str(execution_results)
            investment_strategy_str = str(investment_strategy)
            
            review_result = review_crew.kickoff(
                inputs={
                    "execution_results": execution_results_str,
                    "investment_strategy": investment_strategy_str
                }
            )
            
            # 解析评估结果
            performance = {
                "date": state.get("date", ""),
                "portfolio_value": 0,
                "daily_return": 0,
                "benchmark_return": 0,
                "alpha": 0,
                "sharpe_ratio": 0,
                "max_drawdown": 0,
                "win_rate": 0,
                "improvements": [],
                "summary": str(review_result.raw) if hasattr(review_result, "raw") else str(review_result)
            }
            
            # 返回更新状态
            return {"performance": performance}
        
        except Exception as e:
            logger.error(f"评估投资表现失败: {e}")
            return {"performance": {}, "error": str(e)}


def create_investment_langgraph():
    """创建基于LangGraph的投资决策图"""
    logger.info("创建基于LangGraph的投资决策图")
    
    # 初始化智能体
    market_analyst = MarketAnalyst(verbose=True)
    stock_analyst = StockAnalyst(verbose=True)
    strategy_advisor = StrategyAdvisor(verbose=True)
    portfolio_manager = PortfolioManager(verbose=True)
    chief_investment_officer = ChiefInvestmentOfficer(verbose=True)
    
    # 定义状态类型
    from typing_extensions import TypedDict, NotRequired
    
    class GraphState(TypedDict):
        market_analysis: NotRequired[Dict[str, Any]]
        stock_analysis: NotRequired[Dict[str, Any]]
        investment_strategy: NotRequired[Dict[str, Any]]
        decisions: NotRequired[Dict[str, Any]]
        execution: NotRequired[Dict[str, Any]]
        final_output: NotRequired[str]
    
    # 使用纯字典格式工具，无需工具执行器
    
    # 定义节点函数
    def analyze_market(state: GraphState) -> GraphState:
        """市场分析节点"""
        result = market_analyst._analyze_market_trend()
        return {"market_analysis": {"summary":result}}
    
    def analyze_stocks(state: GraphState) -> GraphState:
        """个股分析节点"""
        # 使用市场分析结果来分析个股
        market_analysis = state.get("market_analysis", {})
        result = {}
        
        for stock_code in ["600519", "000858", "600036"]:
            result[stock_code] = stock_analyst._analyze_stock_fundamentals(stock_code)
        
        return {"stock_analysis": result}
    
    def develop_strategy(state: GraphState) -> GraphState:
        """策略制定节点"""
        market_analysis = state.get("market_analysis", {})
        stock_analysis = state.get("stock_analysis", {})
        
        result = strategy_advisor._develop_investment_strategy(
            market_environment=market_analysis
        )
        
        return {"investment_strategy": result}
    
    def make_decisions(state: GraphState) -> GraphState:
        """决策节点"""
        investment_strategy = state.get("investment_strategy", {})
        stock_analysis = state.get("stock_analysis", {})
        
        result = chief_investment_officer._make_decision(
            analysis_reports=stock_analysis,
            market_conditions={"strategy": investment_strategy}
        )
        
        return {"decisions": result}
    
    def execute_trades(state: GraphState) -> GraphState:
        """交易执行节点"""
        decisions = state.get("decisions", {})
        
        # 执行交易决策
        trades = decisions.get("actions", [])
        result = portfolio_manager._execute_trades(trades)
        
        return {"execution": result}
    
    def summarize(state: GraphState) -> GraphState:
        """总结节点"""
        market_analysis = state.get("market_analysis", {})
        decisions = state.get("decisions", {})
        execution = state.get("execution", {})
        
        # 生成最终输出
        summary = f"""
        投资决策流程总结:
        
        市场分析: {market_analysis}
        
        投资决策: {decisions}
        
        交易执行: {execution}
        """
        
        return {"final_output": summary}
    
    # 定义状态图
    graph = StateGraph(GraphState)
    
    # 添加节点
    graph.add_node("analyze_market", analyze_market)
    graph.add_node("analyze_stocks", analyze_stocks)
    graph.add_node("develop_strategy", develop_strategy)
    graph.add_node("make_decisions", make_decisions)
    graph.add_node("execute_trades", execute_trades)
    graph.add_node("summarize", summarize)
    
    # 定义边
    graph.add_edge("analyze_market", "analyze_stocks")
    graph.add_edge("analyze_stocks", "develop_strategy")
    graph.add_edge("develop_strategy", "make_decisions")
    graph.add_edge("make_decisions", "execute_trades")
    graph.add_edge("execute_trades", "summarize")
    
    # 设置入口点
    graph.set_entry_point("analyze_market")
    
    # 编译图
    investment_graph = graph.compile()
    
    return investment_graph